import Foundation
import Combine

extension Document: AgentThreadStore {
    func readThreadModel() async -> ThreadModel {
        await store.readAsync().thread
    }

    func modifyThreadModel<ReturnVal>(_ callback: @escaping (inout ThreadModel) -> ReturnVal) async -> ReturnVal {
        await store.modifyAsync { state in
            return callback(&state.thread)
        }
    }

    func threadModelPublisher() -> AnyPublisher<ThreadModel, Never> {
        store.publisher.map(\.thread).eraseToAnyPublisher()
    }
}

extension Document {
    func pause() {
        store.modify { state in
            if state.thread.status == .running {
                state.thread.status = .paused
            }
        }
    }

    func unpause() {
        store.modify { state in
            if state.thread.status == .paused {
                state.thread.status = .running // causes checkCancelOrPause to finish
            }
        }
    }

    func stop() {
        // Cancel the current agent task
        currentAgentTask?.cancel()
        currentAgentTask = nil
        toolModalToPresent = nil

        // Reset typing state
        store.modify { state in
            state.thread.status = .none
            state.thread.cancelCount += 1
        }
    }
    
    func clear() {
        // TODO: Make sure this works when paused
        stop()
        store.modify { state in
            state.thread = .init(cancelCount: state.thread.cancelCount + 1)
            state.terminalVisible = false
        }
        terminal = nil
    }
    
    /// Sends a message to the agent and handles the response
    func send(text: String, attachments: [ContextItem]) async {
        let msg = TaggedLLMMessage(role: .user, content: [.text(text)] + attachments)
        let folderURL = store.model.folder
        let curFile = store.model.selectedFileInEditorRelativeToFolder

        stop()
 
        currentAgentTask = Task {
            guard let llm = try? LLMs.smartAgentModel() else {
                await Alerts.showAppAlert(title: "No API Key", message: "Add your API key in Nat → Settings")
                return
            }
            do {
                let tools: [Tool] = [
                    FileReaderTool(), FileEditorTool(), CodeSearchTool(), FileTreeTool(),
                    TerminalTool(), WebResearchTool(), DeleteFileTool(), GrepTool(),
                    BasicContextTool(document: self, currentFilenameFromXcode: curFile),
                ]
                try await send(message: msg, llm: llm, document: self, tools: tools, folderURL: folderURL, maxIterations: store.model.maxIterations)
            } catch {
                if Task.isCancelled { return }
                // Do nothing (We already handle it)
            }
            if !Task.isCancelled {
                currentAgentTask = nil
            }
        }
    }
}
